import lancedb
import pandas as pd
import numpy as np
from typing import List, Dict

class LocalLanceDBExample:
    """本地 LanceDB 使用示例"""

    def __init__(self, db_path: str = "./lancedb_local"):
        """
        初始化本地 LanceDB 数据库

        Args:
            db_path: 本地数据库路径
        """
        self.db = lancedb.connect(db_path)
        print(f"Connected to local LanceDB at {db_path}")

    def create_sample_data(self) -> List[Dict]:
        """创建示例数据"""
        # 生成一些示例向量数据
        np.random.seed(42)
        data = []
        for i in range(100):
            vector = np.random.randn(128).tolist()  # 128维向量
            data.append({
                "id": i,
                "vector": vector,
                "text": f"Sample text item {i}",
                "category": f"category_{i % 5}",
                "score": float(np.random.rand() * 100)
            })
        return data

    def setup_database(self):
        """设置数据库和表"""
        # 创建示例数据
        sample_data = self.create_sample_data()

        # 创建表
        table_name = "embeddings_table_01"
        self.table = self.db.create_table(table_name, data=sample_data)
        print(f"Created table '{table_name}' with {len(sample_data)} records")

        return self.table

    def vector_search_example(self):
        """向量搜索示例"""
        # 生成查询向量
        query_vector = np.random.randn(128).tolist()

        # 执行向量搜索
        results = self.table.search(query_vector).limit(5).to_pandas()
        print("\n=== Vector Search Results ===")
        print(results[["id", "text", "score"]])

        return results

    def filter_search_example(self):
        """条件过滤搜索示例"""
        # 基于条件的搜索
        results = self.table.search().where("category = 'category_2' AND score > 50").limit(10).to_pandas()
        print("\n=== Filtered Search Results ===")
        print(results[["id", "text", "category", "score"]])

        return results

    def insert_new_data(self):
        """插入新数据示例"""
        new_data = [
            {
                "id": 100,
                "vector": np.random.randn(128).tolist(),
                "text": "New inserted item",
                "category": "category_3",
                "score": 95.5
            }
        ]

        self.table.add(new_data)
        print(f"\nInserted {len(new_data)} new records")

    def get_table_info(self):
        """获取表信息"""
        print(f"\n=== Table Information ===")
        print(f"Table name: {self.table.name}")
        print(f"Row count: {self.table.count_rows()}")
        print(f"Schema: {self.table.schema}")

def main():
    """主函数演示"""
    # 创建本地 LanceDB 实例
    db_example = LocalLanceDBExample("/Users/lifangyu/Documents/project/bigdata/bigdata-project/lance/src/main/resources/lance_data/my_lancedb")

    # 设置数据库
    db_example.setup_database()

    # 显示表信息
    db_example.get_table_info()

    # 向量搜索示例
    db_example.vector_search_example()

    # 条件过滤搜索示例
    db_example.filter_search_example()

    # 插入新数据
    db_example.insert_new_data()

    # 再次显示更新后的表信息
    db_example.get_table_info()

if __name__ == "__main__":
    main()
